/**
 * Vanna.AI Frontend Implementation using jQuery
 * This file contains the jQuery-based implementation of the Vanna.AI frontend
 */

// Application state
const state = {
  messageLog: [],
  suggestedQuestions: null,
  trainingData: null,
  questionHistory: [],
  thinking: false,
  questionAsked: false,
  markedCorrect: null,
  currentView: 'chat',
  currentQuestionId: null
};

// Initialize the application when the document is ready
$(document).ready(function() {
  // Setup event listeners
  setupEventListeners();
  
  // Check URL hash for routing
  const hash = window.location.hash.slice(1);
  if (hash === 'training-data') {
    showTrainingData();
  } else if (hash) {
    // If hash is a question ID, load that question
    loadQuestion(hash);
  } else {
    // Otherwise, show the chat view
    showChatView();
    // Generate initial questions
    generateQuestions();
  }
  
  // Load question history
  getQuestionHistory();
});

// Setup event listeners for all interactive elements
function setupEventListeners() {
  // Navigation
  $('#training-data-btn').on('click', showTrainingData);
  $('#new-question-btn').on('click', () => {
    showChatView();
    
    // Clear the input field
    $('#question-input').val('').focus();
    
    // Show welcome screen and hide any messages
    $('#welcome-screen').show();
    $('#message-container').empty();
  });
  
  // Chat input
  $('#question-input').on('keydown', function(e) {
    if (e.key === 'Enter') {
      sendQuestion();
    }
  });
  $('#send-button').on('click', sendQuestion);
  
  // Dynamically added elements
  $(document).on('click', '.suggested-question-btn', function() {
    const question = $(this).text();
    askQuestion(question);
  });
  
  $(document).on('click', '.history-question', function() {
    const id = $(this).data('id');
    loadQuestion(id);
  });
  
  $(document).on('click', '.rerun-sql-btn', function() {
    const id = state.currentQuestionId;
    if (id) {
      rerunSQL(id);
    }
  });
  
  $(document).on('click', '.mark-correct-btn', function() {
    markCorrect(true);
  });
  
  $(document).on('click', '.mark-incorrect-btn', function() {
    markCorrect(false);
  });
  
  $(document).on('click', '.edit-sql-btn', function() {
    const sqlElement = $(this).closest('.sql-container').find('.sql-message');
    const currentSQL = sqlElement.text();
    
    // Replace with textarea
    sqlElement.replaceWith(`<textarea class="sql-editor p-4 w-full bg-gray-100 border-gray-200 rounded-md text-sm font-mono">${currentSQL}</textarea>`);
    
    // Replace buttons
    const buttonContainer = $(this).closest('.button-container');
    buttonContainer.html(`
      <button class="save-sql-btn py-2 px-3 inline-flex justify-center items-center gap-2 rounded-md border font-medium bg-green-600 text-white hover:bg-green-500">
        Save
      </button>
      <button class="cancel-sql-btn py-2 px-3 inline-flex justify-center items-center gap-2 rounded-md border font-medium bg-white text-gray-700">
        Cancel
      </button>
    `);
  });
  
  $(document).on('click', '.save-sql-btn', function() {
    const sqlEditor = $(this).closest('.message').find('.sql-editor');
    const newSQL = sqlEditor.val();
    
    // Update SQL in the UI
    sqlEditor.replaceWith(`<pre class="sql-message">${newSQL}</pre>`);
    
    // Replace buttons
    const buttonContainer = $(this).closest('.button-container');
    buttonContainer.html(`
      <button class="edit-sql-btn py-2 px-3 inline-flex justify-center items-center gap-2 rounded-md border font-medium bg-white text-gray-700">
        Edit
      </button>
      <button class="rerun-sql-btn py-2 px-3 inline-flex justify-center items-center gap-2 rounded-md border font-medium bg-blue-600 text-white">
        Run SQL
      </button>
    `);
    
    // Update SQL in the backend
    updateSQL(newSQL);
  });
  
  $(document).on('click', '.cancel-sql-btn', function() {
    const sqlEditor = $(this).closest('.message').find('.sql-editor');
    const sqlMessage = state.messageLog.find(msg => msg.type === 'sql');
    const currentSQL = sqlMessage ? sqlMessage.text : '';
    
    // Revert to original SQL
    sqlEditor.replaceWith(`<pre class="sql-message">${currentSQL}</pre>`);
    
    // Replace buttons
    const buttonContainer = $(this).closest('.button-container');
    buttonContainer.html(`
      <button class="edit-sql-btn py-2 px-3 inline-flex justify-center items-center gap-2 rounded-md border font-medium bg-white text-gray-700">
        Edit
      </button>
      <button class="rerun-sql-btn py-2 px-3 inline-flex justify-center items-center gap-2 rounded-md border font-medium bg-blue-600 text-white">
        Run SQL
      </button>
    `);
  });
  
  // Training data page events
  $(document).on('click', '.remove-btn', function() {
    const id = $(this).data('id');
    removeTrainingData(id);
  });
  
  $(document).on('click', '.paginate-prev', function() {
    // Pagination logic for training data
    // Implement if needed
  });
  
  $(document).on('click', '.paginate-next', function() {
    // Pagination logic for training data
    // Implement if needed
  });
  
  $(document).on('submit', '#training-form', function(e) {
    e.preventDefault();
    const formData = {
      question: $('#train-question').val(),
      sql: $('#train-sql').val(),
      ddl: $('#train-ddl').val(),
      documentation: $('#train-documentation').val()
    };
    
    trainModel(formData);
  });
}

// View Management Functions

// Show the chat view
function showChatView() {
  state.currentView = 'chat';
  window.location.hash = '';
  
  // Hide/show appropriate sections
  $('#training-data-area').addClass('hidden');
  $('#chat-area').removeClass('hidden');
  
  // Clear messages if needed
  clearMessages();
  
  // Focus on the question input after a short delay
  setTimeout(() => {
    $('#question-input').focus();
  }, 100);
}

// Show the training data view
function showTrainingData() {
  state.currentView = 'training-data';
  window.location.hash = 'training-data';
  
  // Hide/show appropriate sections
  $('#chat-area').addClass('hidden');
  $('#training-data-area').removeClass('hidden');
  
  // Load training data
  getTrainingData();
}

// API Interaction Functions

// Generate initial suggested questions
function generateQuestions() {
  apiCall('generate_questions', 'GET', {})
    .then(response => {
      if (response.type === 'question_list') {
        state.suggestedQuestions = response.questions;
        renderSuggestedQuestions(response.questions);
      }
    });
}

// Ask a question
function askQuestion(question) {
  // Clear previous state
  clearMessages();
  
  // Add user question to message log
  addMessage({
    type: 'user_question',
    question: question
  });
  
  // Set state
  state.questionAsked = true;
  state.thinking = true;
  
  // Generate SQL
  apiCall('generate_sql', 'GET', { question: question })
    .then(response => {
      addMessage(response);
      
      if (response.type === 'sql') {
        state.currentQuestionId = response.id;
        window.location.hash = response.id;
        
        // Run the SQL
        return apiCall('run_sql', 'GET', { id: response.id });
      }
    })
    .then(response => {
      if (response && response.type === 'df') {
        addMessage(response);
        
        // Generate visualization
        return apiCall('generate_plotly_figure', 'GET', { id: state.currentQuestionId });
      }
    })
    .then(response => {
      if (response && response.type === 'plotly_figure') {
        addMessage(response);
        
        // Update question history
        getQuestionHistory();
        
        // Generate followup questions
        return apiCall('generate_followup_questions', 'GET', { id: response.id });
      }
    })
    .then(response => {
      if (response && response.type === 'question_list') {
        state.suggestedQuestions = response.questions;
        renderSuggestedQuestions(response.questions);
      }
      
      state.thinking = false;
    })
    .catch(error => {
      console.error('Error:', error);
      addMessage({
        type: 'error',
        error: String(error)
      });
      
      state.thinking = false;
    });
}

// Send the current question from the input field
function sendQuestion() {
  const questionInput = $('#question-input');
  const question = questionInput.val().trim();
  
  if (question) {
    askQuestion(question);
    questionInput.val('');
  }
}

// Rerun SQL for a question
function rerunSQL(id) {
  addMessage({
    type: 'user_question',
    question: 'Re-run the SQL'
  });
  
  state.thinking = true;
  
  apiCall('run_sql', 'GET', { id: id })
    .then(response => {
      addMessage(response);
      
      if (response.type === 'df') {
        return apiCall('generate_plotly_figure', 'GET', { id: id });
      }
    })
    .then(response => {
      if (response && response.type === 'plotly_figure') {
        addMessage(response);
        return apiCall('generate_followup_questions', 'GET', { id: id });
      }
    })
    .then(response => {
      if (response && response.type === 'question_list') {
        state.suggestedQuestions = response.questions;
        renderSuggestedQuestions(response.questions);
      }
      
      state.thinking = false;
    })
    .catch(error => {
      console.error('Error:', error);
      addMessage({
        type: 'error',
        error: String(error)
      });
      
      state.thinking = false;
    });
}

// Load a previously asked question by ID
function loadQuestion(id) {
  // Clear UI
  clearMessages();
  
  // Set state
  state.currentView = 'chat';
  state.currentQuestionId = id;
  state.questionAsked = true;
  state.thinking = true;
  
  // Show chat view
  $('#training-data-area').addClass('hidden');
  $('#chat-area').removeClass('hidden');
  
  // Fetch question data
  apiCall('load_question', 'GET', { id: id })
    .then(response => {
      if (response.type === 'question_cache') {
        // Add the question and SQL to the UI
        addMessage({
          type: 'user_question',
          question: response.question
        });
        
        addMessage({
          type: 'sql',
          text: response.sql,
          id: id
        });
        
        // If a dataframe is included
        if (response.df) {
          const dfData = JSON.parse(response.df);
          addMessage({
            type: 'df',
            df: JSON.stringify(dfData)
          });
        }
        
        // If a visualization is included
        if (response.fig) {
          addMessage({
            type: 'plotly_figure',
            fig: response.fig,
            id: id
          });
        }
        
        // If followup questions are included
        if (response.followup_questions) {
          state.suggestedQuestions = response.followup_questions;
          renderSuggestedQuestions(response.followup_questions);
        }
      }
      
      state.thinking = false;
    })
    .catch(error => {
      console.error('Error loading question:', error);
      addMessage({
        type: 'error',
        error: 'Failed to load question: ' + String(error)
      });
      
      state.thinking = false;
    });
}

// Get question history
function getQuestionHistory() {
  apiCall('get_question_history', 'GET', {limit: 100, all: true})
    .then(response => {
      if (response.type === 'question_history') {
        state.questionHistory = response.questions;
        renderQuestionHistory(response.questions);
      }
    });
}

// Get training data
function getTrainingData() {
  apiCall('get_training_data', 'GET', {})
    .then(response => {
      state.trainingData = response;
      renderTrainingData(response);
    });
}

// Remove training data
function removeTrainingData(id) {
  // Show confirmation dialog
  if (confirm('Are you sure you want to remove this training data? This action cannot be undone.')) {
    // Show loading state on the button
    const button = $(`.remove-btn[data-id="${id}"]`);
    const originalText = button.text();
    button.text('Removing...').prop('disabled', true).css('opacity', '0.7');
    
    apiCall('remove_training_data', 'POST', { id: id })
      .then(response => {
        if (response.success) {
          // Show success message and refresh data
          alert('Training data removed successfully');
          getTrainingData();
        } else {
          // Show error and reset button
          alert('Error: ' + (response.error || 'Failed to remove training data'));
          button.text(originalText).prop('disabled', false).css('opacity', '1');
        }
      })
      .catch(error => {
        // Show error and reset button
        console.error('Error removing training data:', error);
        alert('Error: ' + error.message);
        button.text(originalText).prop('disabled', false).css('opacity', '1');
      });
  }
}

// Train the model with new data
function trainModel(data) {
  // Check if at least one field is filled
  if (!data.question && !data.sql && !data.ddl && !data.documentation) {
    alert('Please fill at least one field to add training data.');
    return;
  }

  apiCall('train', 'POST', data)
    .then(response => {
      if (response.id) {
        // Show success message
        addMessage({
          type: 'success',
          message: 'Training data added successfully!'
        });
        
        // Refresh training data view
        getTrainingData();
        
        // Clear form
        $('#train-question').val('');
        $('#train-sql').val('');
        $('#train-ddl').val('');
        $('#train-documentation').val('');
      } else if (response.type === 'error') {
        alert('Error: ' + response.error);
      }
    });
}

// Mark a question as correct or incorrect
function markCorrect(isCorrect) {
  state.markedCorrect = isCorrect;
  
  if (isCorrect) {
    // Get current question and SQL
    const userQuestion = state.messageLog.find(msg => msg.type === 'user_question');
    const sqlMessage = state.messageLog.find(msg => msg.type === 'sql');
    
    if (userQuestion && sqlMessage) {
      const trainingData = {
        question: userQuestion.question,
        sql: sqlMessage.text
      };
      
      // Train the model with this data
      trainModel(trainingData);
      
      // Show confirmation
      addMessage({
        type: 'system_message',
        text: 'Thank you for the feedback! This example has been added to the training data.'
      });
    }
  } else {
    // Add a text input for the user to correct the SQL
    const sqlMessage = state.messageLog.find(msg => msg.type === 'sql');
    if (sqlMessage) {
      addMessage({
        type: 'user_sql',
        sql: sqlMessage.text
      });
    }
  }
}

// Update SQL (when user edits it)
function updateSQL(sql) {
  // Find SQL message in the log and update it
  const sqlIndex = state.messageLog.findIndex(msg => msg.type === 'sql');
  if (sqlIndex !== -1) {
    state.messageLog[sqlIndex].text = sql;
  }
}

// Generic API call function
async function apiCall(endpoint, method, params) {
  try {
    state.thinking = true;
    updateThinkingState();
    
    let url = '';
    let response;
    
    if (method === 'GET') {
      // Build query string
      const queryParams = Object.entries(params)
        .filter(([key, value]) => key !== 'endpoint' && key !== 'addMessage')
        .map(([key, value]) => `${encodeURIComponent(key)}=${encodeURIComponent(value)}`)
        .join('&');
      
      url = `/api/v0/${endpoint}?${queryParams}`;
      response = await fetch(url);
    } else {
      url = `/api/v0/${endpoint}`;
      response = await fetch(url, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json'
        },
        body: JSON.stringify(params)
      });
    }
    
    if (!response.ok) {
      throw new Error('The server returned an error. See the server logs for more details.');
    }
    
    const data = await response.json();
    state.thinking = false;
    updateThinkingState();
    
    return data;
  } catch (error) {
    state.thinking = false;
    updateThinkingState();
    return {
      type: 'error',
      error: String(error)
    };
  }
}

// UI Rendering Functions

// Add a message to the UI
function addMessage(message) {
  // Add to state
  state.messageLog.push(message);
  
  // Render in UI
  const msgContainer = $('#message-container');
  
  // Different rendering based on message type
  switch (message.type) {
    case 'user_question':
      msgContainer.append(`
        <div class="message message-user">
          <div class="flex items-center gap-x-2 mb-2">
            <span class="font-medium text-gray-800 dark:text-white">You</span>
          </div>
          <p>${message.question}</p>
        </div>
      `);
      break;
      
    case 'sql':
      msgContainer.append(`
        <div class="message message-agent">
          <div class="flex items-center gap-x-2 mb-2">
            <span class="font-medium text-gray-800 dark:text-white">Vanna</span>
          </div>
          <p>I've created this SQL query based on your question:</p>
          <div class="sql-container mt-2">
            <pre class="sql-message">${message.text}</pre>
            <div class="button-container flex gap-2 mt-2">
              <button class="edit-sql-btn py-2 px-3 inline-flex justify-center items-center gap-2 rounded-md border font-medium bg-white text-gray-700 dark:bg-slate-800 dark:text-gray-400 dark:border-gray-700">
                Edit
              </button>
              <button class="rerun-sql-btn py-2 px-3 inline-flex justify-center items-center gap-2 rounded-md border font-medium bg-blue-600 text-white hover:bg-blue-500">
                Run SQL
              </button>
            </div>
          </div>
          <div class="feedback mt-2">
            <p class="text-sm text-gray-600 dark:text-gray-400">Was this SQL correct?</p>
            <div class="flex gap-2 mt-1">
              <button class="mark-correct-btn py-1 px-2 inline-flex justify-center items-center gap-2 rounded-md border border-green-200 text-sm font-medium bg-white text-green-600 hover:border-green-500">
                👍 Yes
              </button>
              <button class="mark-incorrect-btn py-1 px-2 inline-flex justify-center items-center gap-2 rounded-md border border-red-200 text-sm font-medium bg-white text-red-500 hover:border-red-500">
                👎 No
              </button>
            </div>
          </div>
        </div>
      `);
      break;
      
    case 'df':
      // Parse dataframe JSON
      const dfData = JSON.parse(message.df);
      
      // Create table header
      let tableHeader = '<tr>';
      if (dfData.length > 0) {
        Object.keys(dfData[0]).forEach(key => {
          tableHeader += `<th>${key}</th>`;
        });
      }
      tableHeader += '</tr>';
      
      // Create table rows
      let tableRows = '';
      dfData.forEach(row => {
        tableRows += '<tr>';
        Object.values(row).forEach(value => {
          tableRows += `<td>${value}</td>`;
        });
        tableRows += '</tr>';
      });
      
      msgContainer.append(`
        <div class="message message-agent">
          <div class="flex items-center gap-x-2 mb-2">
            <span class="font-medium text-gray-800 dark:text-white">Vanna</span>
          </div>
          <p>Here are the results from running the SQL query:</p>
          <div class="table-container mt-2">
            <table class="data-table">
              <thead>${tableHeader}</thead>
              <tbody>${tableRows}</tbody>
            </table>
          </div>
          <div class="mt-2">
            <a href="/api/v0/download_csv?id=${state.currentQuestionId}" class="py-2 px-3 inline-flex justify-center items-center gap-2 rounded-md border font-medium bg-white text-gray-700 hover:bg-gray-50 dark:bg-slate-800 dark:text-gray-400 dark:border-gray-700 dark:hover:bg-slate-700">
              Download CSV
            </a>
          </div>
        </div>
      `);
      break;
      
    case 'plotly_figure':
      // Create a container for the plot
      const plotId = `plot-${Date.now()}`;
      
      msgContainer.append(`
        <div class="message message-agent">
          <div class="flex items-center gap-x-2 mb-2">
            <span class="font-medium text-gray-800 dark:text-white">Vanna</span>
          </div>
          <p>Here's a visualization of the data:</p>
          <div id="${plotId}" class="mt-2" style="width:100%; height:400px;"></div>
        </div>
      `);
      
      // Render the Plotly figure
      const figJson = typeof message.fig === 'string' ? JSON.parse(message.fig) : message.fig;
      Plotly.newPlot(plotId, figJson.data, figJson.layout);
      break;
      
    case 'user_sql':
      msgContainer.append(`
        <div class="message message-user">
          <div class="flex items-center gap-x-2 mb-2">
            <span class="font-medium text-gray-800 dark:text-white">You</span>
          </div>
          <p>I'd like to correct the SQL:</p>
          <div class="mt-2">
            <textarea id="corrected-sql" class="p-4 w-full bg-gray-100 border-gray-200 rounded-md text-sm font-mono dark:bg-slate-800 dark:border-gray-700 dark:text-gray-400">${message.sql}</textarea>
            <div class="flex gap-2 mt-2">
              <button id="submit-corrected-sql" class="py-2 px-3 inline-flex justify-center items-center gap-2 rounded-md border font-medium bg-blue-600 text-white hover:bg-blue-500">
                Submit Corrected SQL
              </button>
            </div>
          </div>
        </div>
      `);
      
      // Add event listener for the submit button
      $('#submit-corrected-sql').on('click', function() {
        const correctedSQL = $('#corrected-sql').val();
        
        // Train model with corrected SQL
        const userQuestion = state.messageLog.find(msg => msg.type === 'user_question');
        if (userQuestion) {
          trainModel({
            question: userQuestion.question,
            sql: correctedSQL
          });
          
          // Remove the user SQL message
          state.messageLog = state.messageLog.filter(msg => msg.type !== 'user_sql');
          
          // Add a confirmation message
          addMessage({
            type: 'system_message',
            text: 'Thank you for the correction! This example has been added to the training data.'
          });
        }
      });
      break;
      
    case 'system_message':
      msgContainer.append(`
        <div class="message message-system">
          <p class="text-gray-600 dark:text-gray-400 text-center">${message.text}</p>
        </div>
      `);
      break;
      
    case 'error':
      msgContainer.append(`
        <div class="message message-error bg-red-50 text-red-700 dark:bg-red-900/20 dark:text-red-400">
          <div class="flex items-center gap-x-2 mb-2">
            <span class="font-medium">Error</span>
          </div>
          <p>${message.error}</p>
        </div>
      `);
      break;
    
    case 'success':
      msgContainer.append(`
        <div class="message message-agent">
          <div class="message-header">
            <div class="message-avatar">
              <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="currentColor">
                <path fill-rule="evenodd" d="M2.25 12c0-5.385 4.365-9.75 9.75-9.75s9.75 4.365 9.75 9.75-4.365 9.75-9.75 9.75S2.25 17.385 2.25 12zm13.36-1.814a.75.75 0 10-1.22-.872l-3.236 4.53L9.53 12.22a.75.75 0 00-1.06 1.06l2.25 2.25a.75.75 0 001.14-.094l3.75-5.25z" clip-rule="evenodd" />
              </svg>
            </div>
            <span class="message-name">Success</span>
          </div>
          <p>${message.message}</p>
        </div>
      `);
      break;
  }
  
  // Scroll to bottom
  scrollToBottom();
  
  // Update thinking state
  updateThinkingState();
}

// Render suggested questions
function renderSuggestedQuestions(questions) {
  const suggestionContainer = $('#suggested-questions');
  
  if (questions && questions.length > 0) {
    suggestionContainer.empty().removeClass('hidden');
    
    questions.forEach(question => {
      suggestionContainer.append(`
        <button class="suggested-question-btn">
          ${question}
        </button>
      `);
    });
  } else {
    suggestionContainer.addClass('hidden');
  }
}

// Render question history in the sidebar
function renderQuestionHistory(questions) {
  const historyContainer = $('#question-history');
  historyContainer.empty();
  
  if (questions && questions.length > 0) {
    // Sort questions to show most recent first
    questions.sort((a, b) => {
      // If the ID contains timestamps, use that; otherwise use the alphanumeric sort
      const idA = a.id.split('-').pop();
      const idB = b.id.split('-').pop();
      return idB.localeCompare(idA);
    });

    // Display all questions without limiting the number
    questions.forEach(item => {
      const isActive = item.id === state.currentQuestionId;
      const activeClass = isActive ? 'active' : '';
      
      // Truncate question text more aggressively to fit more items
      const truncatedQuestion = item.question.length > 20 ? 
        item.question.substring(0, 17) + '...' : 
        item.question;
      
      historyContainer.append(`
        <button class="history-question sidebar-nav-item ${activeClass}" data-id="${item.id}" title="${item.question}">
          <svg xmlns="http://www.w3.org/2000/svg" width="14" height="14" fill="currentColor" viewBox="0 0 16 16">
            <path d="M8 15c4.418 0 8-3.134 8-7s-3.582-7-8-7-8 3.134-8 7c0 1.76.743 3.37 1.97 4.6-.097 1.016-.417 2.13-.771 2.966-.079.186.074.394.273.362 2.256-.37 3.597-.938 4.18-1.234A9.06 9.06 0 0 0 8 15z"/>
          </svg>
          <span class="sidebar-nav-text">${truncatedQuestion}</span>
        </button>
      `);
    });
    
    // If there's an active item, scroll to it
    if (state.currentQuestionId) {
      const activeItem = historyContainer.find(`.history-question.active`);
      if (activeItem.length) {
        setTimeout(() => {
          activeItem[0].scrollIntoView({ behavior: 'auto', block: 'nearest' });
        }, 100);
      }
    }
  }
}

// Render training data
function renderTrainingData(data) {
  const trainingDataContainer = $('#training-data-area');
  
  if (data.type === 'df') {
    // Parse dataframe JSON
    const trainingData = JSON.parse(data.df);
    
    // Build table header
    let tableHeader = '<tr>';
    const columns = trainingData.length > 0 ? Object.keys(trainingData[0]) : [];
    
    // Make sure the ID column is first
    const idColumnIndex = columns.indexOf('id');
    if (idColumnIndex !== -1) {
      columns.splice(idColumnIndex, 1);
      columns.unshift('id');
    }
    
    columns.forEach(col => {
      tableHeader += `<th>${col}</th>`;
    });
    tableHeader += '<th>Actions</th></tr>';
    
    // Build table body
    let tableRows = '';
    trainingData.forEach(row => {
      tableRows += '<tr>';
      columns.forEach(col => {
        tableRows += `<td>${row[col] || ''}</td>`;
      });
      
      // Add remove button
      tableRows += `
        <td>
          <button class="remove-btn" data-id="${row.id}">
            Remove
          </button>
        </td>
      `;
      tableRows += '</tr>';
    });
    
    // Create form for adding new training data with improved styling
    const trainingForm = `
      <div class="training-data-form mt-8">
        <div class="form-card">
          <div class="form-card-header">
            <h2 class="form-card-title">Add New Training Data</h2>
            <p class="form-card-subtitle">Add or remove training data. Good training data is the key to accuracy.</p>
          </div>
          <form id="training-form">
            <div class="form-group">
              <label for="train-question" class="form-label">
                Question <span class="optional-label">(optional)</span>
              </label>
              <input type="text" id="train-question" class="form-control training-data-field" placeholder="Enter a natural language question">
            </div>
            
            <div class="form-group">
              <label for="train-sql" class="form-label">
                SQL <span class="optional-label">(optional)</span>
              </label>
              <textarea id="train-sql" class="form-control training-data-field" rows="4" placeholder="Enter SQL code here"></textarea>
            </div>
            
            <div class="form-group">
              <label for="train-ddl" class="form-label">
                DDL <span class="optional-label">(optional)</span>
              </label>
              <textarea id="train-ddl" class="form-control training-data-field" rows="4" placeholder="Enter database schema definition"></textarea>
            </div>
            
            <div class="form-group">
              <label for="train-documentation" class="form-label">
                Documentation <span class="optional-label">(optional)</span>
              </label>
              <textarea id="train-documentation" class="form-control training-data-field" rows="4" placeholder="Enter any additional documentation"></textarea>
            </div>
            
            <div class="form-actions">
              <button type="submit" class="btn-submit">
                Add Training Data
              </button>
            </div>
          </form>
        </div>
      </div>
    `;
    
    // Render the full training data page
    trainingDataContainer.html(`
      <div class="max-w-[85rem] px-4 py-10 sm:px-6 lg:px-8 lg:py-14 mx-auto">
        <div class="training-data-section">
          <h2 class="training-data-section-title">Training Data</h2>
          <p class="training-data-section-description">Manage your training examples for Vanna.AI</p>
          
          <div class="bg-white border border-gray-200 rounded-xl shadow-sm overflow-hidden dark:bg-slate-900 dark:border-gray-700">
            <div class="overflow-x-auto">
              <table class="training-data-table min-w-full">
                <thead>
                  ${tableHeader}
                </thead>
                <tbody>
                  ${tableRows}
                </tbody>
              </table>
            </div>
          </div>
        </div>
        ${trainingForm}
      </div>
    `);
  } else if (data.type === 'error') {
    trainingDataContainer.html(`
      <div class="max-w-4xl mx-auto py-10">
        <div class="bg-red-50 border border-red-200 text-red-700 p-4 rounded-lg">
          <h3 class="font-bold">Error</h3>
          <p>${data.error}</p>
        </div>
      </div>
    `);
  }
}

// Utility Functions

// Clear all messages in the UI
function clearMessages() {
  // Reset state
  state.messageLog = [];
  state.markedCorrect = null;
  
  // Clear UI
  $('#message-container').empty();
  $('#welcome-screen').show();
}

// Update the thinking state in the UI
function updateThinkingState() {
  if (state.thinking) {
    // Remove welcome screen if there are messages
    if (state.messageLog.length > 0) {
      $('#welcome-screen').hide();
    }
    
    // Add thinking indicator if not already present
    if ($('#thinking-indicator').length === 0) {
      $('#message-container').append(`
        <div id="thinking-indicator" class="message message-agent">
          <div class="flex items-center gap-x-2">
            <img src="/vanna.svg" class="w-8 h-8 animate-bounce" alt="Vanna thinking">
            <span class="font-medium text-gray-800 dark:text-white">Vanna is thinking...</span>
          </div>
        </div>
      `);
      
      scrollToBottom();
    }
  } else {
    // Remove thinking indicator
    $('#thinking-indicator').remove();
  }
}

// Scroll to the bottom of the message container
function scrollToBottom() {
  const messageContainer = document.getElementById('message-container');
  if (messageContainer) {
    messageContainer.scrollTop = messageContainer.scrollHeight;
  }
} 